# mapping_with_model_robust_v2.py
# 功能：增加了题号格式规范化功能，兼容图片识别的题号与JSON文件中的题号格式差异。

import os
import json
import base64
import datetime
from openai import OpenAI
import re # 正则表达式模块，用于格式规范化

# ====== 路径配置（请根据实际情况修改为你的绝对路径） ======
BASE_DIR = r"E:\沙粒云\自媒体\2025视频制作\20250610ai高考\流程\anwser\data"
IMAGE_FOLDER = r"E:\沙粒云\自媒体\2025视频制作\20250610ai高考\流程\aigaokao_ui\pic"
ANWSER_DIR = os.path.join(BASE_DIR, "anwser")
ANSWER_JSON_PATH = os.path.join(ANWSER_DIR, "answer_list.json")
OUTPUT_MAPPING_PATH = os.path.join(ANWSER_DIR, "mapping_result_robust_v2.txt")
LOG_FILE = os.path.join(ANWSER_DIR, "mapping_model_robust_v2_log.txt")

# ====== Qwen‐VL 接口配置 ======
API_KEY = "sk-c48712427ba84b59b9c89fc2883e48b6"  # 请替换为你真实的 API_KEY
BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
QWEN_VL_MODEL = "qwen-vl-max"

# ====== 初始化 OpenAI 客户端 & 日志文件 ======
qwen_client = OpenAI(api_key=API_KEY, base_url=BASE_URL)

with open(LOG_FILE, "w", encoding="utf-8") as f:
    f.write(f"[{datetime.datetime.now()}] 鲁棒映射任务日志文件(v2)创建/重置\n")

def log(msg: str):
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    line = f"[{timestamp}] {msg}"
    print(line)
    with open(LOG_FILE, "a", encoding="utf-8") as f:
        f.write(line + "\n")

# --- 其他函数保持不变 (encode_image_to_base64, get_image_mime_type) ---
def encode_image_to_base64(path: str) -> str:
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

def get_image_mime_type(path: str) -> str:
    ext = os.path.splitext(path)[1].lower()
    if ext in ['.jpg', '.jpeg']: return 'image/jpeg'
    if ext == '.png': return 'image/png'
    if ext == '.webp': return 'image/webp'
    raise ValueError(f"Unsupported image format: {ext}")
# ---

## [新增功能] 题号格式规范化函数
def normalize_question_id(text: str) -> str:
    """
    将各种格式的题号字符串（如 '第14题', '14.', '13. (1)'）
    统一转换为纯净的内部ID（如 '14', '13(1)'）。
    """
    # 查找所有数字和英文括号
    parts = re.findall(r'\d+|\(|\)', text)
    return "".join(parts)

def load_flat_answers(answer_json_path: str) -> list:
    with open(answer_json_path, "r", encoding="utf-8") as f:
        units = json.load(f)
    flat_list = []
    for unit in units:
        # 注意：这里需要处理正文可能不在 "answers" 列表中的情况
        for item in unit.get("answers", []):
            if isinstance(item, str) and ":" in item:
                parts = item.replace("：", ":").split(":", 1)
                q = parts[0].strip()
                a = parts[1].strip()
                flat_list.append({"question": q, "answer": a})
            else:
                # 对于像 "正文：" 这样的条目，或者没有冒号的，可以选择忽略或特殊处理
                # 这里我们简单忽略非 “题号:答案” 格式的字符串
                pass 
                
        # 针对您JSON中作文题的特殊处理
        if unit.get("unit") == "Unit7":
            title_item = unit["answers"][0] # "第23题: 标题：问渠那得清如许"
            content_items = unit["answers"][1:] # 剩余部分是正文
            
            # 提取标题
            q_title = title_item.split(":")[0].strip()
            a_title = title_item.split(":", 1)[1].strip()
            flat_list.append({"question": q_title, "answer": a_title})
            
            # 拼接正文 (这里我们假设作文题只有一个，即第23题的正文)
            # 我们将正文内容归属于同一个题号，例如 "第23题-正文"
            # 为了简单起见，这里可以将整个作文合并为一个答案
            full_content = "\n".join(content_items)
            # 假设作文内容也属于第23题，可以创建一个新的题号或附加到原题上
            # 这里我们简单地不再为“正文”创建新条目，因为映射通常是针对选择题、填空题
            # 如果作文也有红框，需要定义其题号，如 "第23题-正文"
            pass

    return flat_list


def build_page_prompt(page_num: int, image_path: str, image_b64: str) -> list:
    mime_type = get_image_mime_type(image_path)
    messages = [
        {
            "role": "system",
            "content": (
                "你是一个精准的试卷映射机器人。你的任务是识别图像中的红色答题框，并将其与试卷上印刷的题号关联起来。"
                "你必须严格按照用户要求的 JSON 格式输出，不要添加任何解释或无关文字。"
            )
        },
        {
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{image_b64}"}},
                {
                    "type": "text",
                    "text": (
                        f"这是第 {page_num} 页的扫描图像。\n\n"
                        "你的任务是：\n"
                        "1. 按“从上到下、从左到右”的顺序，识别图像中的每一个红色方框。\n"
                        "2. 对于每一个红框，请准确识别它旁边印刷的题号（例如 '12.', '13.(1)' 等）。这个题号是关键的映射依据。\n"
                        "3. 输出一个 JSON 数组，其中每个对象包含三个字段：'question' (从图片中识别的题号), 'page' (当前页码), 和 'box_idx' (从1开始的红框序号)。\n\n"
                        "输出格式必须严格如下所示（这是一个例子）：\n"
                        "```json\n"
                        "[\n"
                        '  {"question": "12.", "page": 2, "box_idx": 1},\n'
                        '  {"question": "13.(1)", "page": 2, "box_idx": 2},\n'
                        '  {"question": "13.(2)", "page": 2, "box_idx": 3}\n'
                        "]\n"
                        "```\n\n"
                        "重要提醒：\n"
                        f"- 当前页码是: {page_num}。\n"
                        "- 如果本页没有任何红色方框，请只返回一个空的 JSON 数组 `[]`。\n"
                        "- 不要输出任何 JSON 代码块标记（```json ... ```）之外的任何文字。\n"
                        "- `question` 字段的值必须与你在图片上看到的题号完全一致。"
                    )
                }
            ]
        }
    ]
    return messages

def main():
    log("====== 鲁棒映射脚本(v2)启动 ======")

    flat_answers = load_flat_answers(ANSWER_JSON_PATH)
    
    ## [修改点 1] 创建答案查找字典时，使用规范化后的题号作为 KEY
    # 同时，VALUE存储一个包含原始题号和答案的字典，方便后续使用
    answer_lookup = {
        normalize_question_id(item["question"]): {
            "original_q": item["question"], 
            "answer": item["answer"]
        } 
        for item in flat_answers
    }
    total_questions = len(flat_answers)
    log(f"共加载到 {total_questions} 条答案记录。创建了 {len(answer_lookup)} 条规范化索引。")

    try:
        img_files_raw = [f for f in os.listdir(IMAGE_FOLDER) if f.lower().endswith((".jpg", ".jpeg"))]
        img_files = sorted(img_files_raw, key=lambda f: int(os.path.splitext(f)[0]))
    except (FileNotFoundError, ValueError) as e:
        log(f"[ERROR] 图像文件夹或文件名不规范：{IMAGE_FOLDER}。错误：{e}")
        return
    
    if not img_files:
        log(f"[ERROR] 在 {IMAGE_FOLDER} 中未找到任何图像文件，程序退出。")
        return

    final_mapping_list = []

    for page_idx, img_name in enumerate(img_files, start=1):
        img_path = os.path.join(IMAGE_FOLDER, img_name)
        log(f"\n—— 处理第 {page_idx} 页 图像：{img_name} ——")

        try:
            image_b64 = encode_image_to_base64(img_path)
        except Exception as e:
            log(f"[ERROR] 无法读取或编码图片 {img_name}：{e}，跳过该页。")
            continue

        messages = build_page_prompt(page_idx, img_path, image_b64)

        try:
            log(f"Step: 向 Qwen‐VL 提交第 {page_idx} 页映射请求…")
            response = qwen_client.chat.completions.create(
                model=QWEN_VL_MODEL,
                messages=messages,
                stream=False
            )
            model_output = response.choices[0].message.content.strip()
            log(f"[MODEL OUTPUT]\n{model_output}")

        except Exception as e:
            log(f"[ERROR] Qwen‐VL 请求失败：{e}，跳过第 {page_idx} 页。")
            continue

        try:
            json_match = re.search(r'\[.*\]', model_output, re.DOTALL)
            if not json_match:
                log(f"[WARN] 在模型输出中未找到有效的 JSON 数组。跳过第 {page_idx} 页。")
                continue
            
            page_mappings = json.loads(json_match.group(0))

            if not page_mappings:
                log(f"第 {page_idx} 页没有红框（模型返回空列表），继续下一页。")
                continue

            page_mapped_count = 0
            for item in page_mappings:
                if not all(k in item for k in ['question', 'page', 'box_idx']):
                    log(f"[WARN] JSON 对象格式错误，缺少键: {item}，跳过。")
                    continue

                # [修改点 2] 对从AI获取的题号也进行规范化，然后用规范化ID进行查找
                model_question_id = item['question']
                normalized_id = normalize_question_id(model_question_id)
                
                answer_data = answer_lookup.get(normalized_id)
                
                if answer_data is None:
                    # 这条日志现在只会在真正找不到题号时触发（比如AI识别错了数字）
                    log(f"[WARN] 规范化后的题号 '{normalized_id}' (来自AI识别的'{model_question_id}') 在答案列表中未找到，跳过该条记录。")
                    continue
                
                final_mapping_list.append({
                    "question": answer_data["original_q"], # 存入原始的、带“第X题”的题号
                    "answer": answer_data["answer"],
                    "page": item['page'],
                    "box_idx": item['box_idx']
                })
                page_mapped_count += 1
            
            log(f"第 {page_idx} 页成功解析并映射 {page_mapped_count} 条记录。")

        except Exception as e:
            log(f"[ERROR] 处理第 {page_idx} 页的映射结果时发生未知错误: {e}")

    log(f"\n====== 所有页面处理完毕，共映射到 {len(final_mapping_list)} 条记录（原始答案总计 {total_questions} 条）。 ======")
    with open(OUTPUT_MAPPING_PATH, "w", encoding="utf-8") as fout:
        final_mapping_list.sort(key=lambda x: (x['page'], x['box_idx']))
        for entry in final_mapping_list:
            line_to_write = f"{entry['page']}-{entry['box_idx']}：{entry['answer']}\n"
            fout.write(line_to_write)
    log(f"映射表已保存到：{OUTPUT_MAPPING_PATH}")

    print("\n====== 映射表预览 (新格式) ======")
    for entry in final_mapping_list:
        print(f"{entry['page']}-{entry['box_idx']}：{entry['answer']}")

    log("====== 脚本结束 ======")

if __name__ == "__main__":
    main()